#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import torch
import torch.nn.functional as F
import os
import struct

def im2col_torch(x, kernel_size, stride=1, padding=0):
    """
    PyTorch实现的im2col，使用F.unfold
    x: [N, C, H, W] tensor
    return: [N, C*kH*kW, L] where L = out_H * out_W
    """
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    
    # 使用PyTorch的unfold函数
    col = F.unfold(x, kernel_size, stride=stride, padding=padding)
    return col

def gen_golden_data_im2col_torch():
    input_shape = (1, 1, 9, 257)
    input_data_np = np.random.uniform(low=-1.0, high=1.0, size=input_shape).astype(np.float32)
    input_data = torch.from_numpy(input_data_np)
    
    # im2col参数
    kernel_h = 2
    kernel_w = 2
    kernel_size = (kernel_h, kernel_w)
    stride_val = 1           # 步长
    padding_val = 0          # 填充
    
    # 计算im2col
    golden_tensor = im2col_torch(input_data, kernel_size, stride_val, padding_val)
    golden = golden_tensor.numpy().astype(np.float32)
    
    # 创建目录
    os.makedirs("input", exist_ok=True)
    os.makedirs("output", exist_ok=True)
    
    # 根据 Im2ColCustomTilingData 结构体生成tiling数据
    # struct Im2ColCustomTilingData {
    #     uint32_t totalLength;  // N * C * H * W
    #     uint32_t tileNum;
    # 
    #     // 参数设置
    #     int32_t N;
    #     int32_t C;
    #     int32_t H;
    #     int32_t W;
    # 
    #     // 填充参数
    #     int32_t kernel_h;
    #     int32_t kernel_w;
    #     int32_t stride_val;
    #     int32_t padding_val;
    # };
    
    # 获取输入形状参数
    N, C, H, W = input_shape
    total_length = N * C * H * W  # totalLength
    tile_num = 8  # 假设使用8个tile
    
    # 使用struct.pack确保正确的内存布局
    # 'I' = unsigned int (uint32_t), 'i' = int (int32_t)
    tiling_data = struct.pack('IIiiiiiiii', 
                             total_length,  # uint32_t totalLength
                             tile_num,      # uint32_t tileNum
                             N,             # int32_t N
                             C,             # int32_t C  
                             H,             # int32_t H
                             W,             # int32_t W
                             kernel_h,      # int32_t kernel_h
                             kernel_w,      # int32_t kernel_w
                             stride_val,    # int32_t stride_val
                             padding_val)   # int32_t padding_val
    
    # 保存tiling数据
    with open("input/input_tiling.bin", "wb") as f:
        f.write(tiling_data)
    
    # 保存输入和输出数据
    input_data_np.tofile("input/input_x.bin")
    golden.tofile("output/golden.bin")

if __name__ == "__main__":
    # 使用PyTorch实现
    gen_golden_data_im2col_torch()